from qutip import *
import numpy as np
sys.path.append('D:\Dropbox\Dokumente\PI3\Code\Python\quantum-tools')
import qutip_enhanced

qte = qutip_enhanced.QutipEnhanced()

bloch = Bloch3d()
rho0_e = ket2dm(basis(2, 0))
rho0_n = ket2dm(basis(2, 0))
rho = tensor(rho0_e, rho0_n)

def round_odd(n):
    return 2*round((n+1)/2.) - 1

def n(hf, target_tau):
    return 2*target_tau*hf

def tau2(hf, target_tau):
    return round_odd(n(hf, target_tau))/(2*hf)

def rel_tau(hf, target_tau):
    return abs(tau(hf,target_tau) - target_tau)/target_tau

target_tau = 120
hf = 0.013
a = 2*np.pi*hf/2.
tau = tau2(hf, target_tau)
f_rabi_n = 2*np.pi/(2.*tau)
b = 0.5*f_rabi_n
h = Qobj([[-a, 0, 0, 0],
          [0,  a, 0, 0],
          [0,  0, 0, b],
          [0,  0, b, 0]])
h.dims = [[2, 2], [2, 2]]
rho = qte.rotate(rho, rotation_axis={'y': -1}, angle=np.pi/2., rotated_spin=0)
rho = mesolve(h, rho, np.linspace(0, tau, 2), [], []).states[-1]
rho = qte.rotate(rho, rotation_axis={'y': -1}, angle=np.pi, rotated_spin=0)
rho = mesolve(h, rho, np.linspace(0, tau, 2), [], []).states[-1]
bloch.add_states(rho.ptrace(0))
bloch.show()